import random
import copy

from typing import Union, List
from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig,CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_aclnn_dequant_rope_quant_kvcache")
class DefaultGenerator(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)
        self.x_shape = []
        self.cos_shape = []
        self.D = None
        self.Nq = None
        self.Nkv = None
        self.H = None
        self.kCacheRef_shape = []


    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        # x，shape为(B,S,H)或(B,H)，H=D*(Nq+Nkv+Nkv)，H<=4096且64对齐
        if index == 0:
            D_range_value = []
            for i in range(64, 1365, 64):
                D_range_value.append(i)
            self.D = random.choice(D_range_value)
            max_nq_2nkv = (4095 // self.D)
            max_nq = max_nq_2nkv - 2
            self.Nq = random.randint(1, max_nq)
            max_nkv = (max_nq_2nkv - self.Nq) // 2
            self.Nkv = random.randint(1, max_nkv)
            self.H = self.D * (self.Nq + 2 * self.Nkv)
            # if self.H >= 8192:
            #     print("self.H:", self.H)
            # print("111111111111111111")
            input_case.shape[-1] = self.H

            tol=input_case.shape[0]*input_case.shape[1]
            if len(input_case.shape)==3:
                tol=tol*input_case.shape[2]
            while tol>11000000:
                input_case.shape[0]=random.randint(1, 15)
                if len(input_case.shape)==3:
                    input_case.shape[1]=random.randint(1, 15)
                tol=input_case.shape[0]*input_case.shape[1]
                if len(input_case.shape)==3:
                    tol=tol*input_case.shape[2]
            
            self.x_shape = input_case.shape
            # print(input_case)



        ## cos，shape为(B,S,1,D)或(B,D)，x是3维时cos是4维，x是2维时cos是2维
        if index == 1:
            input_case.shape = []
            # print("222222222222222222222")
            if len(self.x_shape) == 2:
                input_case.shape.append(self.x_shape[0])
                input_case.shape.append(self.D)
            elif len(self.x_shape) == 3:
                input_case.shape = self.x_shape[:2]
                input_case.shape.append(1)
                input_case.shape.append(self.D)
            self.cos_shape = input_case.shape
            # print(input_case)

        ## sin，shape为(B,S,1,D)或(B,D)，x是3维时cos是4维，x是2维时cos是2维
        if index == 2:
            # print("33333333333333333333")
            input_case.shape = self.cos_shape
            # print(input_case)

        ## kCacheRef，shape是(C_1,C_2,Nkv,D)
        if index == 3:
            # print("4444444444444444444")
            input_case.shape[2] = self.Nkv
            input_case.shape[3] = self.D
            self.kCacheRef_shape = input_case.shape
            # print(input_case)

        ## vCacheRef，shape是(C_1,C_2,Nkv,D)
        if index == 4:
            # print("555555555555555555555555")
            input_case.shape = self.kCacheRef_shape
            # print(input_case)

        # ## indices，shape是(B*S)或(B)
        # if index == 5:
        #     # print("66666666666666666666")
        #     input_case.shape[0] = self.x_shape[0]
        #     # print(input_case)

        ## scaleK，shape是(D)
        if index == 6:
            # print("7777777777777777777777")
            input_case.shape[0] = self.Nkv * self.D
            # print(input_case)

        ## scaleV，shape是(D)
        if index == 7:
            # print("888888888888888888888")
            input_case.shape[0] = self.Nkv * self.D
            # print(input_case)

        ## offsetKOptional，shape是(D)
        if index == 8:
            # print("99999999999999999")
            input_case.shape[0] = self.Nkv * self.D
            # print(input_case)

        ## offsetVOptional，shape是(D)
        if index == 9:
            # print("10101010101010101010")
            input_case.shape[0] = self.Nkv * self.D
            # print(input_case)

        ## weightScaleOptional，shape是(H)
        if index == 10:
            # print("11111")
            input_case.shape[0] = self.H
            # print(input_case)

        ## activationScaleOptional，shape是(B*S)或(B)
        if index == 11:
            # print("12121212121212121212")
            if len(self.x_shape) == 2:
                input_case.shape[0] = self.x_shape[0]
            elif len(self.x_shape) == 3:
                input_case.shape[0] = self.x_shape[0]*self.x_shape[1]
            # print(input_case)

        ## biasOptional，shape是(H)
        if index == 12:
            # print("13131313131313131313")
            input_case.shape[0] = self.H
            # print(input_case)

        ## sizeSplits，取值为[D*Nq,D*Nkv,D*Nkv]
        if index == 13:
            # print("141414141414141414")
            input_case[0].range_values = self.D*self.Nq
            input_case[1].range_values = self.D*self.Nkv
            input_case[2].range_values = self.D*self.Nkv
            # print(input_case)

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        # print("----------------------------------")
        # print(case_config.inputs)

        # cache_mode为"contiguous"时，C1>B，C2>S,indices的shape为(B)，取值大于等于0小于等于C2-S，x的shape为(B,H)时S=1
        if case_config.inputs[-1].range_values == "contiguous":
            case_config.inputs[4].shape[0] = case_config.inputs[0].shape[0] + random.randint(1, 50)
            case_config.inputs[4].shape[1] = case_config.inputs[0].shape[1] + random.randint(1, 50)
            case_config.inputs[3].shape[0] = case_config.inputs[4].shape[0]
            case_config.inputs[3].shape[1] = case_config.inputs[4].shape[1]
            if len(case_config.inputs[0].shape) == 2:
                max_indices = case_config.inputs[4].shape[1] - 1
            else:
                max_indices = case_config.inputs[4].shape[1] - case_config.inputs[0].shape[1]
            case_config.inputs[5].range_values = [0,max_indices]
            case_config.inputs[5].shape[0] = case_config.inputs[0].shape[0]
        else:
            # cache_mode为"page"时，C1*C2>B*S,indices的shape为(B*S)，取值大于等于0小于C1*C2且不重复，x的shape为(B,H)时S=1
            C1xC2 = case_config.inputs[4].shape[0]*case_config.inputs[4].shape[1]
            if len(case_config.inputs[0].shape) == 2:
                case_config.inputs[4].shape[1] = ((case_config.inputs[0].shape[0] + random.randint(1, 50))
                                                  // case_config.inputs[4].shape[0]) + 1
                case_config.inputs[5].shape[0] = case_config.inputs[0].shape[0]
            else:
                case_config.inputs[4].shape[1] = ((case_config.inputs[0].shape[0]*case_config.inputs[0].shape[1]
                                                   + random.randint(1, 50)) // case_config.inputs[4].shape[0]) + 1
                case_config.inputs[5].shape[0] = case_config.inputs[0].shape[0]*case_config.inputs[0].shape[1]
            case_config.inputs[3].shape[1] = case_config.inputs[4].shape[1]
            # print("----------------------------------")
            # print(case_config.inputs)
            max_indices = case_config.inputs[4].shape[0]*case_config.inputs[4].shape[1] - 1
            case_config.inputs[5].range_values = [0,max_indices]


        # x不为int32时，x、cos、sin的dtype保持一致
        # x为int32时，cos、sin的dtype保持一致
        if case_config.inputs[0].dtype != "int32":
            case_config.inputs[1].dtype = case_config.inputs[0].dtype
            case_config.inputs[2].dtype = case_config.inputs[0].dtype
        else:
            case_config.inputs[2].dtype = case_config.inputs[1].dtype

        # print(case_config.inputs)
        return case_config